package com.smallstep.apilimit.ratelimiter;

import com.alibaba.fastjson.JSON;
import com.google.common.base.Preconditions;
import com.google.common.math.LongMath;
import org.apache.commons.lang3.StringUtils;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.concurrent.TimeUnit;

/**
 * @author 李俊
 * @Description
 * @Date 2019/9/4 16:52
 */
public class RedisRateLimiter {
    /**
     * key Redis key
     */
    private String key;
    /**
     * permitsPerSecond 每秒放入的令牌数
     */
    private Double permitsPerSecond;
    /**
     * maxBurstSeconds 最大存储maxBurstSeconds秒生成的令牌
     */
    private Integer maxBurstSeconds = 60;

    private SyncLock syncLock;



    private StringRedisTemplate stringRedisTemplate;

    public static RedisRateLimiter create(String key, Double permitsPerSecond,
                                          StringRedisTemplate stringRedisTemplate) {
       return new RedisRateLimiter(key,permitsPerSecond,1,stringRedisTemplate);
    }


    public RedisRateLimiter(String key, Double permitsPerSecond,
                            Integer maxBurstSeconds, StringRedisTemplate stringRedisTemplate) {
        this.key = key;
        this.maxBurstSeconds = maxBurstSeconds;
        this.permitsPerSecond = permitsPerSecond;
        syncLock = new SyncLock(key+":RateLimitLock",10L,50L,stringRedisTemplate);
        this.stringRedisTemplate = stringRedisTemplate;
    }

    /**
     * 获取redis内的值
     * @return
     */
    RedisPerimits getRedisPerimits() {
        RedisPerimits redisPerimits = null;
       String redisPerimitsStr = stringRedisTemplate.opsForValue().get(key+":RateLimit");
       if (StringUtils.isEmpty(redisPerimitsStr)) {
           redisPerimits = new RedisPerimits(permitsPerSecond,maxBurstSeconds,
                   now());
           setRedisPerimits(redisPerimits);
       } else {
           redisPerimits = JSON.parseObject(redisPerimitsStr,RedisPerimits.class);
       }
       return redisPerimits;
    }

    /**
     *
     * @param redisPerimits
     */
    void setRedisPerimits(RedisPerimits redisPerimits){
        String redisPerimitsStr = JSON.toJSONString(redisPerimits);
        stringRedisTemplate.opsForValue()
                .set(key+":RateLimit",redisPerimitsStr,
                        redisPerimits.expires(),TimeUnit.SECONDS);
    }


    Long acquire(Long tokens) {
        Long milliToWait = reserve(tokens);
        try {
            Thread.sleep(milliToWait);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return milliToWait;
    }

    Long acquire() {
        return acquire(1L);
    }

    private Long reserve(Long tokens) {
        checkTokens(tokens);
        try {
            syncLock.lock();
            return reserveAndGetWaitLength(tokens);
        } finally {
            syncLock.unLock();
        }
    }

    /**
     * 尝试获取token
     * @param tokens
     * @param timeout
     * @param unit
     * @return
     */
    Boolean tryAcquire(Long tokens, Long timeout, TimeUnit unit) {
        Long timeoutMicros = Math.max(unit.toMillis(timeout), 0);
        checkTokens(tokens);

        Long milliToWait;
        try {
            syncLock.lock();
            if (!canAcquire(tokens, timeoutMicros)) {
                return false;
            } else {
                milliToWait = reserveAndGetWaitLength(tokens);
            }
        } finally {
            syncLock.unLock();
        }
        try {
            Thread.sleep(milliToWait);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return true;
    }

    private Boolean  canAcquire(Long tokens, Long timeoutMillis){
        return queryEarliestAvailable(tokens) - timeoutMillis <= 0;
    }

    private Long queryEarliestAvailable(Long tokens) {
        Long n = now();
        RedisPerimits redisPerimits = getRedisPerimits();
        redisPerimits.reSync(n);
        Long storedPermitsToSpend = Math.min(tokens, redisPerimits.storedPermits); // 可以消耗的令牌数
        Long freshPermits = tokens - storedPermitsToSpend; // 需要等待的令牌数
        Long waitMillis = freshPermits * redisPerimits.intervalMillis; // 需要等待的时间
        return LongMath.saturatedAdd(redisPerimits.nextFreeTicketMillis - n, waitMillis);
    }

    /**
     *
     * @param tokens
     */
    private void checkTokens(Long tokens) {
        Preconditions.checkArgument(tokens > 0,
                "Requested tokens $tokens must be positive");
    }
    /**
     *
     * @return
     */
    public Boolean tryAcquire() {
      return  tryAcquire(1L, 0L, TimeUnit.MICROSECONDS);
    }


    private Long reserveAndGetWaitLength(Long tokens) {
        Long n = now();
        RedisPerimits redisPerimits = getRedisPerimits();
        redisPerimits.reSync(n);
        Long storedPermitsToSpend = Math.min(tokens, redisPerimits.storedPermits); // 可以消耗的令牌数
        Long freshPermits = tokens - storedPermitsToSpend; // 需要等待的令牌数
        Long waitMillis = freshPermits * redisPerimits.intervalMillis; // 需要等待的时间
        redisPerimits.nextFreeTicketMillis = LongMath.saturatedAdd(redisPerimits.nextFreeTicketMillis, waitMillis);
        redisPerimits.storedPermits -= storedPermitsToSpend;
        setRedisPerimits(redisPerimits);
        return redisPerimits.nextFreeTicketMillis - n;
    }

    public Long now() {
        return System.currentTimeMillis();
    }
}
